package com.ruoyi.common.scoket;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import javax.websocket.server.PathParam;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@Component
public class UploadProgressHandler extends TextWebSocketHandler {

    private final Map<String, PutObjectProgressListener> uploadProgressMap = new ConcurrentHashMap<>();
    private static final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        uploadProgressMap.forEach((requestId, progressListener) -> {
            try {
                session.sendMessage(new TextMessage(progressListener.getProgress()));
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception{
        Object params = session.getAttributes().get("requestId");
        sessionMap.put(params.toString(), session);
        super.afterConnectionEstablished(session);
    }
    public static Map<String, WebSocketSession> getSessionMap() {
        return sessionMap;
    }
    public void addProgressListener(String requestId, PutObjectProgressListener progressListener) {
        uploadProgressMap.put(requestId, progressListener);
    }
    public void removeProgressListener(String requestId) {
        uploadProgressMap.remove(requestId);
    }
}

